# 构建自己的数据

import glob
import os
import random
from torch.utils.data import Dataset

from PIL import Image

import torchvision.transforms as transforms


class ImageDataset(Dataset):
    def __init__(self, root='', transform=None, model='train'):
        self.transform = transforms.Compose(transform)

        self.pathA = root+model+'A/*'
        self.pathB = root+model+'B/*'
        # print('pathA:', self.pathA)

        self.list_A = glob.glob(self.pathA)
        self.list_B = glob.glob(self.pathB)

    def __getitem__(self, index):
        im_pathA = self.list_A[index % len(self.pathA)]
        im_pathB = random.choice(self.list_B)

        im_A = Image.open(im_pathA)
        im_B = Image.open(im_pathB)

        item_A = self.transform(im_A)
        item_B = self.transform(im_B)

        return {'A': item_A, 'B': item_B}

    def __len__(self):
        return max(len(self.list_A), len(self.list_B))


if __name__ == '__main__':
    from torch.utils.data import DataLoader

    root = '../data/'
    transform_ = [transforms.Resize(256, Image.BILINEAR), transforms.ToTensor()]
    im_dataset = ImageDataset(root, transform_, 'train')
    # print(len(im_dataset))
    dataLoader = DataLoader(im_dataset, batch_size=1, shuffle=True, num_workers=1)

    for i, batch in enumerate(dataLoader):
        print(i, batch)
